local newt = reload("newt")
local class = reload("class")
local store = reload("store")
local clear = reload("table.opt.clear")
local search = reload("table.array.traverse")

local temp = other.tcreate(128, 0)
local newa64, dela64 = store(64, 0)
local newh8, delh8 = store(0, 8)


---@class BPlusTreet @b+树
local this = class()

---创建非叶节点
local function new_internal(parent)
    local node = newh8()
    node.keys = newa64()   -- 路由键数组
    node.childs = newa64() -- 子节点指针数组
    node.parent = parent   -- 父节点指针
    return node
end

---创建叶子节点
local function new_leaf(parent)
    local node = newh8()
    node.keys = newa64() -- 数据键数组 { {uid=_, idx=_, inf=_ }, ... }
    node.parent = parent
    node.next_leaf = nil
    return node
end

---找到第一个<=目标值的索引（非严格匹配）
local function first_big_internal_index(self, node, leaf_key)
    local keys = node.keys
    local low = 1
    local high = #keys
    -- 二分查找第一个严格大于key的索引
    while low <= high do
        local mid = (low + high) // 2
        local ret = self._c(leaf_key, keys[mid])
        if -1 == ret then
            high = mid - 1 -- 目标在左半区
        else
            low = mid + 1  -- 目标在右半区
        end
    end

    -- low 是第一个大于key的索引
    return low
end

--- 查找第一个等于目标值的索引（严格匹配）
local function first_equal_internal_index(node, idx)
    local keys = node.keys
    local low, high = 1, #keys

    while low <= high do
        local mid = (low + high) // 2
        if keys[mid].idx == idx then
            high = mid - 1 -- 继续向左搜索第一个出现的位置
        elseif keys[mid].idx < idx then
            low = mid + 1  -- 目标在右半区
        else
            high = mid - 1 -- 目标在左半区
        end
    end

    return low
end

---内部方法：二分查找叶子节点
local function first_big_leaf_index(self, leaf_node, leaf_key)
    local keys = leaf_node.keys
    local low = 1
    local high = #keys

    -- 二分查找第一个严格大于key的索引
    while low <= high do
        local mid = (low + high) // 2
        local ret = self._c(leaf_key, keys[mid])
        if -1 == ret then
            high = mid - 1 -- 目标在左半区
        else
            low = mid + 1  -- 目标在右半区
        end
    end

    -- low 是第一个大于key的索引
    return low
end

---内部方法：二分查找叶子节点中指定uniq值的索引
local function first_equal_leaf_uniq(self, node, uid)
    local keys = node.keys
    for k, v in ipairs(keys) do
        if v.uid == uid then
            return k
        end
    end
end

---内部方法：二分查找叶子节点
local function first_equal_leaf_index(node, idx)
    local keys = node.keys
    local low, high = 1, #keys

    local pos = nil -- 默认第一个节点
    while low <= high do
        local mid = (low + high) // 2
        if keys[mid].idx == idx then
            pos = mid      -- 记录候选位置
            high = mid - 1 -- 继续向左搜索第一个出现的位置
        elseif keys[mid].idx < idx then
            low = mid + 1  -- 目标在右半区
        else
            high = mid - 1 -- 目标在左半区
        end
    end

    return pos
end

---内部方法：分裂叶子节点
local function split_leaf(self, leaf)
    local new_leaf = new_leaf(leaf.parent)
    local split_pos = math.ceil(self._m / 2) -- 按阶数计算分裂点

    -- 迁移右半部分数据（从split_pos+1开始）
    for i = split_pos + 1, #leaf.keys do
        table.insert(new_leaf.keys, leaf.keys[i])
        self._hash[leaf.keys[i].uid] = new_leaf -- 更新哈希表
    end
    for i = #leaf.keys, split_pos + 1, -1 do
        table.remove(leaf.keys)
    end

    -- 设置正确的路由键（左子节点最大键）
    local split_key = leaf.keys[#leaf.keys]

    -- 维护链表顺序（关键！）
    new_leaf.next_leaf = leaf.next_leaf
    leaf.next_leaf = new_leaf

    return new_leaf, split_key
end

---内部方法：分裂内部节点
local function split_internal(self, internal)
    local new_internal = new_internal(internal.parent)
    local split_pos = math.ceil((self._m - 1) / 2) -- 保证左节点有足够元素

    -- 提升左子节点的最大键
    local split_key = internal.keys[split_pos]

    -- 迁移右半部分
    for i = split_pos + 1, #internal.keys do
        table.insert(new_internal.keys, internal.keys[i])
    end
    for i = #internal.keys, split_pos + 1, -1 do
        table.remove(internal.keys)
    end

    -- 迁移子节点
    for i = split_pos + 1, #internal.childs do
        local child = internal.childs[i]
        child.parent = new_internal
        table.insert(new_internal.childs, child)
    end
    for i = #internal.childs, split_pos + 1, -1 do
        table.remove(internal.childs)
    end

    return new_internal, split_key
end

---内部方法：更新父节点
local function update_parent(self, old_node, new_node, split_key)
    local parent = old_node.parent
    if not parent then
        -- 创建新根节点
        local new_root = new_internal(nil)
        table.insert(new_root.keys, split_key)
        table.insert(new_root.childs, old_node)
        table.insert(new_root.childs, new_node)
        old_node.parent = new_root
        new_node.parent = new_root
        self._root = new_root
        return
    end

    -- 插入到父节点
    local pos = first_big_internal_index(self, parent, split_key)
    table.insert(parent.keys, pos, split_key)
    table.insert(parent.childs, pos + 1, new_node)

    -- 递归处理父节点分裂
    if #parent.keys > self._m then
        local new_parent, new_split_key = split_internal(self, parent)
        update_parent(self, parent, new_parent, new_split_key)
    end
end


---内部方法：处理内部节点下溢
local function handle_internal_underflow(self, parent, current_node, idx_in_parent)
    -- 尝试向左兄弟借键
    local left_sib = idx_in_parent > 1 and parent.childs[idx_in_parent - 1]
    if left_sib and #left_sib.keys >= math.ceil(self._m / 2) then
        local borrowed_key = table.remove(left_sib.keys)
        local borrowed_child = table.remove(left_sib.childs)

        -- 更新父节点键
        table.insert(current_node.keys, 1, parent.keys[idx_in_parent - 1])
        parent.keys[idx_in_parent - 1] = borrowed_key

        -- 更新子节点
        table.insert(current_node.childs, 1, borrowed_child)
        borrowed_child.parent = current_node
        return true
    end

    -- 尝试向右兄弟借键
    local right_sib = idx_in_parent < #parent.childs and parent.childs[idx_in_parent + 1]
    if right_sib and #right_sib.keys >= math.ceil(self._m / 2) then
        local borrowed_key = table.remove(right_sib.keys, 1)
        local borrowed_child = table.remove(right_sib.childs, 1)

        -- 更新父节点键
        table.insert(current_node.keys, parent.keys[idx_in_parent])
        parent.keys[idx_in_parent] = borrowed_key

        -- 更新子节点
        table.insert(current_node.childs, borrowed_child)
        borrowed_child.parent = current_node
        return true
    end

    -- 合并节点
    local merged = false
    if left_sib then
        -- 合并到左兄弟
        table.insert(left_sib.keys, parent.keys[idx_in_parent - 1])
        for _, key in ipairs(current_node.keys) do
            table.insert(left_sib.keys, key)
        end
        for _, child in ipairs(current_node.childs) do
            table.insert(left_sib.childs, child)
            child.parent = left_sib
        end
        table.remove(parent.keys, idx_in_parent - 1)
        table.remove(parent.childs, idx_in_parent)
        merged = true
    elseif right_sib then
        -- 合并到右兄弟
        table.insert(right_sib.keys, 1, parent.keys[idx_in_parent])
        for i = #current_node.keys, 1, -1 do
            table.insert(right_sib.keys, 1, current_node.keys[i])
        end
        for i = #current_node.childs, 1, -1 do
            table.insert(right_sib.childs, 1, current_node.childs[i])
            current_node.childs[i].parent = right_sib
        end
        table.remove(parent.keys, idx_in_parent)
        table.remove(parent.childs, idx_in_parent)
        merged = true
    end

    -- 递归处理父节点
    if merged and parent.parent then
        local grandparent = parent.parent
        local parent_idx = search(grandparent.childs, parent)
        if parent_idx then
            handle_internal_underflow(self, grandparent, parent, parent_idx)
        end
    end

    -- 更新根节点
    if parent == self._root and #parent.keys == 0 then
        self._root = parent.childs[1] or nil
        if self._root then
            self._root.parent = nil
        end
    end
end

---构造函数
function this:ctor(m, sortCompare)
    self:clear()
    self._m = m or 64
    ---用于排序比较
    self._c = sortCompare
end

function this:clear()
    self._root = nil
    self._size = 0
    self._hash = newt(0, 512) --唯一索引
    self._mcnt = newt(0, 512) --索引数量
end

---获取总数量
function this:size()
    return self._size
end

---获取索引数量
function this:idxsum(idx)
    return self._mcnt[idx] or 0
end

---获取数量map
function this:idxMap()
    return self._mcnt
end

---插入数据（核心逻辑）
function this:insert(uid, idx, inf)
    ---重复判断
    local leaf = self._hash[uid]
    if leaf then
        return false
    end

    ---数量统计
    self._size = self._size + 1
    local ocnt = self._mcnt[idx] or 0
    self._mcnt[idx] = ocnt + 1

    local leaf_key = newh8()
    leaf_key.uid = uid
    leaf_key.idx = idx
    leaf_key.inf = inf

    if not self._root then
        self._root = new_leaf(nil)
        table.insert(self._root.keys, leaf_key)
        self._hash[uid] = self._root
        return true
    end

    -- 查找插入路径
    local node = self._root
    while node.childs do
        local pos = first_big_internal_index(self, node, leaf_key)
        node = node.childs[pos]
    end

    -- 插入到叶子节点
    local pos = first_big_leaf_index(self, node, leaf_key)
    table.insert(node.keys, pos, leaf_key)
    self._hash[uid] = node

    -- 处理分裂
    while #node.keys > self._m do
        local new_leaf = split_leaf(self, node)
        local split_key = new_leaf.keys[1]
        if node.parent then
            update_parent(self, node, new_leaf, split_key)
            node = node.parent
        else
            update_parent(self, node, new_leaf, split_key)
            break
        end
    end
    return true
end

---删除数据（完整实现）
function this:delete(uid)
    if not self._root then
        return false
    end

    -- 通过哈希获取目标叶子节点
    local leaf = self._hash[uid]
    if not leaf then
        return false
    end

    ---数量统计
    self._size = self._size - 1
    self._hash[uid] = nil

    -- 查找键在叶子中的位置
    local deleted_pos = first_equal_leaf_uniq(self, leaf, uid)
    if not deleted_pos then
        assert(false, "删除的键不存在 uid:" .. uid)
        return false
    end

    ---类型统计
    local dkey = leaf.keys[deleted_pos]
    local ocnt = self._mcnt[dkey.idx] or 0
    self._mcnt[dkey.idx] = ocnt - 1
    table.remove(leaf.keys, deleted_pos)

    -- 处理叶子节点下溢
    if #leaf.keys < math.ceil(self._m / 2) and leaf.parent then
        local parent = leaf.parent
        local parent_idx = search(parent.childs, leaf)
        if not parent_idx then
            return true
        end

        -- 尝试向左兄弟借键
        local left_sib = parent_idx > 1 and parent.childs[parent_idx - 1]
        if left_sib and #left_sib.keys >= math.ceil(self._m / 2) then
            local borrowed = table.remove(left_sib.keys)
            table.insert(leaf.keys, 1, borrowed)
            parent.keys[parent_idx - 1] = left_sib.keys[#left_sib.keys]
            self._hash[borrowed.uid] = leaf
            return true
        end

        -- 尝试向右兄弟借键
        local right_sib = parent_idx < #parent.childs and parent.childs[parent_idx + 1]
        if right_sib and #right_sib.keys >= math.ceil(self._m / 2) then
            local borrowed = table.remove(right_sib.keys, 1)
            table.insert(leaf.keys, borrowed)
            parent.keys[parent_idx] = right_sib.keys[1]
            self._hash[borrowed.uid] = leaf
            return true
        end

        -- 合并节点
        if left_sib then
            -- 合并到左兄弟
            for _, key in ipairs(leaf.keys) do
                table.insert(left_sib.keys, key)
                self._hash[key.uid] = left_sib
            end
            left_sib.next_leaf = leaf.next_leaf
            table.remove(parent.childs, parent_idx)
            table.remove(parent.keys, parent_idx - 1)
        elseif right_sib then
            -- 合并到右兄弟
            for i = #leaf.keys, 1, -1 do
                table.insert(right_sib.keys, 1, leaf.keys[i])
                self._hash[leaf.keys[i].uid] = right_sib
            end
            right_sib.next_leaf = leaf.next_leaf
            table.remove(parent.childs, parent_idx)
            table.remove(parent.keys, parent_idx)

            -- 强制更新父节点路由键（关键修复）
            if parent_idx <= #parent.keys then
                parent.keys[parent_idx] = right_sib.keys[1]
            elseif parent_idx > 1 then
                parent.keys[parent_idx - 1] = right_sib.keys[1]
            end
        end

        -- 处理父节点下溢
        if parent.parent or (#parent.keys < math.ceil(self._m / 2) and parent ~= self._root) then
            local grandparent = parent.parent
            local parent_idx = grandparent and search(grandparent.childs, parent)
            if parent_idx then
                handle_internal_underflow(self, grandparent, parent, parent_idx)
            end
        end

        -- 根节点处理
        if parent == self._root and #parent.keys == 0 then
            self._root = parent.childs[1] or nil
            if self._root then
                self._root.parent = nil
            end
        end
    elseif not leaf.parent and #leaf.keys == 0 then
        -- 处理根节点为空的情况
        self._root = nil
    end

    delh8(dkey)
    return true
end

---范围查询（完整实现）
function this:range(searIndex, skipCount, needCount)
    skipCount = skipCount or 0
    needCount = needCount or self:idxsum(searIndex)
    local results = clear(temp)
    if not self._root then
        return results
    end

    -- 找到起始叶子节点
    local node = self._root
    while node.childs do
        local idx = first_equal_internal_index(node, searIndex)
        node = node.childs[idx]
    end

    -- 遍历叶子节点
    local count = 0
    while node do
        ---这个节点找不到说明其他节点叶找不到
        local pos = first_equal_leaf_index(node, searIndex)

        local len = #node.keys
        while pos and pos <= len do
            local key = node.keys[pos]
            local skip = false
            if key.idx == searIndex then
                if skipCount > 0 then
                    local leftCount = len - pos
                    if skipCount >= leftCount then
                        skipCount = skipCount - leftCount
                        pos = #node.keys
                    else
                        pos = pos + skipCount
                        skipCount = 0
                    end
                    skip = true
                else
                    table.insert(results, key)
                    count = count + 1
                    if count >= needCount then
                        return results
                    end
                end
            elseif key.idx > searIndex then
                ---后面是找不到了直接退出
                return results
            end

            if not skip then
                pos = pos + 1
            end
        end
        node = node.next_leaf
    end
    return results
end

return this
